#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import socket
import threading
import struct
import argparse

class ModbusTCPBridge:
    def __init__(self, tcp_server_host='127.0.0.1', tcp_server_port=9010):
        self.tcp_server_host = tcp_server_host
        self.tcp_server_port = tcp_server_port
        self.tcp_server_socket = None
        self.modbus_socket = None
        self.modbus_host = None
        self.modbus_port = None
        self.running = True
        
    def connect_to_tcp_server(self):
        """连接到TCP服务器"""
        try:
            self.tcp_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.tcp_server_socket.connect((self.tcp_server_host, self.tcp_server_port))
            print(f"已连接到TCP服务器 {self.tcp_server_host}:{self.tcp_server_port}")
            
            # 启动接收TCP服务器消息的线程
            receive_thread = threading.Thread(target=self.receive_tcp_messages, daemon=True)
            receive_thread.start()
            return True
        except Exception as e:
            print(f"连接TCP服务器失败: {e}")
            return False
    
    def receive_tcp_messages(self):
        """接收TCP服务器消息的线程函数"""
        while self.running and self.tcp_server_socket:
            try:
                data = self.tcp_server_socket.recv(1024)
                if data:
                    try:
                        if data.startswith(b'text:'):
                            message = data[5:].decode('utf-8')
                            print(f"\n[TCP服务器文本] {message}")
                        else:
                            message = data.decode('utf-8')
                            print(f"\n[TCP服务器] {message}")
                    except UnicodeDecodeError:
                        print(f"\n[TCP服务器数据] {data.hex().upper()}")
                    print("控制台> ", end="", flush=True)
                else:
                    print("\nTCP服务器断开连接")
                    break
            except Exception as e:
                if self.running:
                    print(f"\n接收TCP服务器消息错误: {e}")
                break
    
    def connect_to_modbus_device(self, host, port):
        """连接到Modbus TCP设备"""
        try:
            if self.modbus_socket:
                self.modbus_socket.close()
            
            self.modbus_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.modbus_socket.settimeout(10)
            self.modbus_socket.connect((host, port))
            self.modbus_host = host
            self.modbus_port = port
            print(f"已连接到Modbus设备 {host}:{port}")
            return True
        except Exception as e:
            print(f"连接Modbus设备失败: {e}")
            return False
    
    def calculate_crc(self, data):
        """计算Modbus RTU CRC校验码"""
        crc = 0xFFFF
        for pos in data:
            crc ^= pos
            for _ in range(8):
                if (crc & 1) != 0:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc.to_bytes(2, 'little')

    def parse_modbus_command(self, cmd_str):
        """解析Modbus RTU格式指令，如果没有CRC则自动计算并添加"""
        try:
            parts = cmd_str.strip().split()

            # 检查指令长度，支持带CRC和不带CRC的格式
            if len(parts) < 6:
                raise ValueError("指令格式应为: 单元ID 功能码 起始地址H 起始地址L 数量H 数量L [CRCH CRCL]")

            unit_id = int(parts[0], 16)
            function_code = int(parts[1], 16)
            start_addr = (int(parts[2], 16) << 8) + int(parts[3], 16)
            quantity = (int(parts[4], 16) << 8) + int(parts[5], 16)

            # 构建RTU数据部分(不带CRC)
            rtu_data = bytes([unit_id, function_code]) + struct.pack(">HH", start_addr, quantity)

            # 检查是否已经包含CRC
            if len(parts) >= 8:
                # 已包含CRC，验证是否正确
                provided_crc = bytes([int(parts[6], 16), int(parts[7], 16)])
                calculated_crc = self.calculate_crc(rtu_data)

                if provided_crc != calculated_crc:
                    print(f"警告: 提供的CRC ({provided_crc.hex().upper()}) 与计算的CRC ({calculated_crc.hex().upper()}) 不匹配")
                    print(f"使用计算的CRC: {calculated_crc.hex().upper()}")

                crc = calculated_crc
            else:
                # 没有CRC，自动计算
                crc = self.calculate_crc(rtu_data)
                print(f"自动计算CRC: {crc.hex().upper()}")

            return unit_id, function_code, start_addr, quantity, rtu_data + crc

        except ValueError as e:
            raise ValueError(f"无效指令: {e}")
    
    def send_modbus_command(self, unit_id, function_code, start_addr, quantity):
        """发送Modbus TCP命令并返回响应"""
        if not self.modbus_socket:
            print("错误: 未连接到Modbus设备")
            return None
        
        try:
            # 构建Modbus TCP PDU
            transaction_id = 0x0001
            protocol_id = 0x0000
            length = 6
            
            # 打包请求数据
            modbus_data = struct.pack(
                ">HHHBBHH",
                transaction_id,
                protocol_id,
                length,
                unit_id,
                function_code,
                start_addr,
                quantity
            )
            
            # 发送请求
            self.modbus_socket.send(modbus_data)
            print(f"发送Modbus TCP: {modbus_data.hex().upper()}")
            
            # 接收响应
            response = self.modbus_socket.recv(1024)
            print(f"接收Modbus响应: {response.hex().upper()}")
            
            return response
            
        except Exception as e:
            print(f"发送Modbus命令错误: {e}")
            return None
    
    def convert_tcp_to_rtu_response(self, tcp_response):
        """将Modbus TCP响应转换为RTU格式"""
        if not tcp_response or len(tcp_response) < 8:
            return None

        try:
            # 解析Modbus TCP头
            (trans_id, prot_id, length, unit_id, func_code) = struct.unpack(">HHHBB", tcp_response[:8])

            # 提取RTU部分（去掉TCP头）
            rtu_data = tcp_response[6:]  # 从unit_id开始的所有数据

            # 计算CRC并添加
            crc = self.calculate_crc(rtu_data)
            rtu_with_crc = rtu_data + crc

            return rtu_with_crc

        except Exception as e:
            print(f"TCP转RTU错误: {e}")
            return None

    def parse_modbus_response(self, response, function_code):
        """解析Modbus响应数据"""
        if not response or len(response) < 8:
            return None

        try:
            # 解析Modbus TCP头
            (trans_id, prot_id, length, unit_id, func_code) = struct.unpack(">HHHBB", response[:8])

            if func_code == function_code:
                byte_count = response[8]
                data = response[9:9+byte_count]

                # 如果是读取寄存器的响应
                if function_code in [3, 4] and byte_count % 2 == 0:
                    registers = []
                    for i in range(0, byte_count, 2):
                        registers.append((data[i] << 8) + data[i+1])
                    return {
                        'success': True,
                        'data': data.hex().upper(),
                        'registers': registers
                    }
                else:
                    return {
                        'success': True,
                        'data': data.hex().upper()
                    }
            else:
                error_code = response[8] if len(response) > 8 else 0
                return {
                    'success': False,
                    'error': f"功能码={func_code}, 错误码={error_code}"
                }
        except Exception as e:
            print(f"解析响应错误: {e}")
            return None
    
    def send_to_tcp_server(self, data):
        """发送数据到TCP服务器"""
        if not self.tcp_server_socket:
            print("错误: 未连接到TCP服务器")
            return False
        
        try:
            if isinstance(data, str):
                self.tcp_server_socket.send(data.encode('utf-8'))
            else:
                self.tcp_server_socket.send(data)
            return True
        except Exception as e:
            print(f"发送到TCP服务器错误: {e}")
            return False
    
    def run_interactive_mode(self):
        """运行交互模式"""
        print("Modbus TCP 桥接工具")
        print("命令说明:")
        print("  ch [host] [port]     - 连接到Modbus TCP设备")
        print("                         默认: ch 127.0.0.1 5020")
        print("                         示例: ch (使用默认地址)")
        print("                         示例: ch 192.168.1.100 502")
        print("  cl <modbus_cmd>      - 发送Modbus采集指令")
        print("                         格式: cl 单元ID 功能码 地址H 地址L 数量H 数量L [CRCH CRCL]")
        print("                         示例: cl 01 03 00 01 00 01 (CRC自动计算)")
        print("                         示例: cl 01 03 00 00 00 06 (带CRC)")
        print("  send <message>       - 直接发送消息到TCP服务器")
        print("  status               - 显示连接状态")
        print("  quit                 - 退出程序")
        print("  其他任何输入         - 原样发送到TCP服务器")
        print()
        
        while self.running:
            try:
                command = input("控制台> ").strip()
                
                if not command:
                    continue
                
                if command.lower() == 'quit':
                    self.running = False
                    break
                
                elif command.lower() == 'status':
                    tcp_status = "已连接" if self.tcp_server_socket else "未连接"
                    modbus_status = f"已连接到 {self.modbus_host}:{self.modbus_port}" if self.modbus_socket else "未连接"
                    print(f"TCP服务器: {tcp_status}")
                    print(f"Modbus设备: {modbus_status}")
                
                elif command.lower().startswith('ch'):
                    # 连接到Modbus设备
                    parts = command.split()

                    # 默认值
                    default_host = '127.0.0.1'
                    default_port = 5020

                    if len(parts) == 1:
                        # 只输入 "ch"，使用默认值
                        host = default_host
                        port = default_port
                        print(f"使用默认Modbus设备地址: {host}:{port}")
                    elif len(parts) == 2:
                        print("用法: ch [host] [port] 或 ch (使用默认值)")
                        print(f"默认值: {default_host}:{default_port}")
                        continue
                    elif len(parts) >= 3:
                        try:
                            host = parts[1]
                            port = int(parts[2])
                        except ValueError:
                            print("错误: 端口必须是数字")
                            continue

                    self.connect_to_modbus_device(host, port)
                
                elif command.lower().startswith('cl '):
                    # 发送Modbus采集指令
                    modbus_cmd = command[3:].strip()
                    if not modbus_cmd:
                        print("用法: cl <modbus_cmd> (如: cl 01 03 00 01 00 01)")
                        continue

                    try:
                        unit_id, function_code, start_addr, quantity, rtu_with_crc = self.parse_modbus_command(modbus_cmd)
                        print(f"完整RTU指令: {rtu_with_crc.hex().upper()}")

                        response = self.send_modbus_command(unit_id, function_code, start_addr, quantity)

                        if response:
                            # 将Modbus TCP响应转换为RTU格式后发送到TCP服务器
                            rtu_response = self.convert_tcp_to_rtu_response(response)
                            if rtu_response:
                                if self.send_to_tcp_server(rtu_response):
                                    print(f"Modbus RTU数据已发送到TCP服务器: {rtu_response.hex().upper()}")
                                    print(f"原始TCP响应: {response.hex().upper()}")
                                else:
                                    print("发送到TCP服务器失败")
                            else:
                                print("TCP转RTU格式失败")

                            # 同时解析并显示给用户看
                            parsed = self.parse_modbus_response(response, function_code)
                            if parsed:
                                if parsed['success']:
                                    print(f"采集成功: {parsed['data']}")
                                    if 'registers' in parsed:
                                        print(f"寄存器值: {parsed['registers']}")
                                else:
                                    print(f"采集失败: {parsed['error']}")
                            else:
                                print("解析响应失败")
                        else:
                            print("采集失败")

                    except ValueError as e:
                        print(f"指令错误: {e}")

                elif command.lower().startswith('send '):
                    # 直接发送消息到TCP服务器
                    message = command[5:].strip()
                    if message:
                        if self.send_to_tcp_server(message):
                            print("消息已发送到TCP服务器")
                        else:
                            print("发送失败")
                    else:
                        print("用法: send <message>")

                else:
                    # 如果不是已知命令，将输入内容原样发送到TCP服务器
                    if self.send_to_tcp_server(command):
                        print(f"消息已发送到TCP服务器: {command}")
                    else:
                        print("发送到TCP服务器失败")
            
            except KeyboardInterrupt:
                print("\n正在退出...")
                self.running = False
                break
            except Exception as e:
                print(f"命令执行错误: {e}")
    
    def close(self):
        """关闭所有连接"""
        self.running = False
        
        if self.modbus_socket:
            try:
                self.modbus_socket.close()
                print("Modbus连接已关闭")
            except:
                pass
        
        if self.tcp_server_socket:
            try:
                self.tcp_server_socket.close()
                print("TCP服务器连接已关闭")
            except:
                pass

def test_tcp_to_rtu_conversion():
    """测试TCP到RTU转换功能"""
    bridge = ModbusTCPBridge()

    # 测试数据：你提供的TCP响应
    tcp_response = bytes.fromhex("000100000005010302C14B")
    print(f"原始TCP响应: {tcp_response.hex().upper()}")

    # 转换为RTU格式
    rtu_response = bridge.convert_tcp_to_rtu_response(tcp_response)
    if rtu_response:
        print(f"转换后RTU格式: {rtu_response.hex().upper()}")

        # 验证RTU格式的组成
        unit_id = rtu_response[0]
        func_code = rtu_response[1]
        byte_count = rtu_response[2]
        data = rtu_response[3:3+byte_count]
        crc = rtu_response[-2:]

        print(f"单元ID: 0x{unit_id:02X}")
        print(f"功能码: 0x{func_code:02X}")
        print(f"字节数: {byte_count}")
        print(f"数据: {data.hex().upper()}")
        print(f"CRC: {crc.hex().upper()}")
    else:
        print("转换失败")

def bridge_main():
    parser = argparse.ArgumentParser(description='Modbus TCP 桥接工具')
    parser.add_argument('--tcp-host', help='TCP服务器IP地址')
    parser.add_argument('--tcp-port', type=int, help='TCP服务器端口')
    parser.add_argument('--test', action='store_true', help='运行转换测试')
    parser.add_argument('--no-input', action='store_true', help='不提示输入，直接使用默认值')
    args = parser.parse_args()

    if args.test:
        test_tcp_to_rtu_conversion()
        return

    # 默认TCP服务器配置
    default_tcp_host = '127.0.0.1'
    default_tcp_port = 9010

    # 确定TCP服务器地址和端口
    if args.tcp_host and args.tcp_port:
        tcp_host = args.tcp_host
        tcp_port = args.tcp_port
    elif args.no_input:
        tcp_host = args.tcp_host or default_tcp_host
        tcp_port = args.tcp_port or default_tcp_port
    else:
        # 交互式输入TCP服务器地址
        print("=== TCP服务器连接配置 ===")
        tcp_host_input = input(f"请输入TCP服务器IP地址 (默认: {default_tcp_host}): ").strip()
        tcp_host = tcp_host_input if tcp_host_input else default_tcp_host

        tcp_port_input = input(f"请输入TCP服务器端口 (默认: {default_tcp_port}): ").strip()
        if tcp_port_input:
            try:
                tcp_port = int(tcp_port_input)
            except ValueError:
                print("端口输入无效，使用默认端口")
                tcp_port = default_tcp_port
        else:
            tcp_port = default_tcp_port

        print(f"将连接到TCP服务器: {tcp_host}:{tcp_port}")
        print()

    bridge = ModbusTCPBridge(tcp_host, tcp_port)

    try:
        # 连接到TCP服务器
        if not bridge.connect_to_tcp_server():
            print("无法连接到TCP服务器，程序退出")
            return

        # 运行交互模式
        bridge.run_interactive_mode()

    except Exception as e:
        print(f"程序错误: {e}")
    finally:
        bridge.close()

if __name__ == '__main__':
    bridge_main()
